from matplotlib.patches import FancyArrow
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from typing import List, Tuple, Dict, Literal


class Plane2D:
    """在二维平面绘制图像"""

    @staticmethod
    def enable_chinese(fontfamily: str='SimHei'):
        plt.rcParams['font.sans-serif'] = [fontfamily]   #用来正常显示中文标签
        plt.rcParams['axes.unicode_minus'] = False   #用来正常显示负号

    @staticmethod
    def use_tex():
        matplotlib.use('pgf')
        # plt.rcParams['text.usetex'] = True
        # plt.rcParams['text.latex.preamble'] = r'\usepackage{ctex}'

    @staticmethod
    def show():
        plt.show()

    def __init__(self,
                 xrange: Tuple[float, float]=None,
                 yrange: Tuple[float, float]=None,
                 x_expand_factor: float=0.1,
                 y_expand_factor: float=0.1,
                 xticks: List[float]=None,
                 yticks: List[float]=None,
                 xticklabels: List[str]=None,
                 yticklabels: List[str]=None,
                 title: str=None,
                 axis_type: Literal['center', 'box']='box'):
        """初始化坐标轴。

        xrange: x 轴的范围
        yrange: y 轴的范围
        x_expand_factor: x 轴向两端扩张的长度占 xrange 的比例。
        y_expand_factor: y 轴向两端扩张的长度占 yrange 的比例。
        xticks: x 轴刻度
        yticks: y 轴刻度
        xticklabels: x 轴刻度标签
        yticklabels: y 轴刻度标签
        title: 图题
        """
        self._axis_type = axis_type
        self._zorder = 0.0
        self.fig, self.ax = plt.subplots()
        self.ax.axis('equal')
        # self.ax.grid(True)
        if xrange:
            xrange = list(xrange)
            xexpand = (xrange[1] - xrange[0]) * x_expand_factor / 2
            xrange[0] -= xexpand
            xrange[1] += xexpand
            self.ax.set_xlim(*xrange)
        if yrange:
            yrange = list(yrange)
            yexpand = (yrange[1] - yrange[0]) * y_expand_factor / 2
            yrange[0] -= yexpand
            yrange[1] += yexpand
            self.ax.set_ylim(*yrange)
        if xrange and yrange:
            w = 6.4
            h = w * ((yrange[1] - yrange[0]) / (xrange[1] - xrange[0]))
            self.fig.set_size_inches(w, h)
        if xticks:
            self.ax.set_xticks(xticks)
        if xticklabels:
            self.ax.set_xticklabels(xticklabels)
        if yticks:
            self.ax.set_yticks(yticks)
        if yticklabels:
            self.ax.set_yticklabels(yticklabels)
        if title:
            self.ax.set_title(title)
        if self._axis_type == "center":
            self.ax.set_axis_off()
            axis_params = dict(x=0, y=0, dx=0, dy=0,  # to be overwritten
                length_includes_head=True, overhang=.3, color='black', zorder=0)
            self.xaxis = self.ax.arrow(**axis_params)
            self.yaxis = self.ax.arrow(**axis_params)
            self.update_axes()

    def update_axes(self):
        if self._axis_type != "center":
            return
        xmin, xmax, ymin, ymax = self.ax.axis()
        dx = xmax - xmin
        dy = ymax - ymin
        head_width = 0.015 * np.hypot(dx, dy)
        head_length = head_width * 1.5
        self.xaxis: FancyArrow
        self.xaxis.set_data(
            x=xmin, y=0, dx=dx, dy=0,
            head_width=head_width,
            head_length=head_length,
        )
        self.yaxis.set_data(
            x=0, y=ymin, dx=0, dy=dy,
            head_width=head_width,
            head_length=head_length,
        )

    def step_zorder(self):
        self._zorder += 0.01

    def point(self, p: Tuple[float, float], label: str=None):
        """绘制一个点

        x: 横坐标
        y: 纵坐标
        """
        self.step_zorder()
        r = self.ax.scatter(*p, label=label, zorder=self._zorder)
        self.update_axes()
        return r

    def points(self, ps: List[Tuple[float, float]], label: str=None):
        """绘制一个点

        x: 横坐标
        y: 纵坐标
        """
        self.step_zorder()
        r = self.ax.scatter(*zip(*ps), label=label, zorder=self._zorder)
        self.update_axes()
        return r

    def line(self, p1: Tuple[float, float], p2: Tuple[float, float], label: str=None):
        """绘制直线。

        参数：
        p1：第一个点的 (x, y) 坐标
        p2：第二个点的 (x, y) 坐标
        """
        self.step_zorder()
        r = self.ax.plot(*zip(p1, p2), label=label, zorder=self._zorder)
        self.update_axes()
        return r

    def dashline(self, p1: Tuple[float, float], p2: Tuple[float, float], label: str=None):
        """绘制直线。

        参数：
        p1：第一个点的 (x, y) 坐标
        p2：第二个点的 (x, y) 坐标
        """
        self.step_zorder()
        r = self.ax.plot(*zip(p1, p2), label=label, linestyle='--', zorder=self._zorder)
        self.update_axes()
        return r

    def arrow(self, 
              p1: Tuple[float, float],
              p2: Tuple[float, float],
              linewidth: float=None,
              head_width: float=None,
              head_length: float=None,
              overhang: float=0):
        """绘制箭头 p1 -> p2。

        参数：
        p1：第一个点的 (x, y) 坐标
        p2：第二个点的 (x, y) 坐标
        """
        self.step_zorder()
        p1 = np.array(p1)
        p2 = np.array(p2)
        dp = p2 - p1
        if not head_width:
            if head_length:
                head_width = head_length / 1.5
            else:
                xlim = np.array(self.ax.get_xlim())
                ylim = np.array(self.ax.get_ylim())
                xlim[0] = min([xlim[0], p1[0], p2[0]])
                xlim[1] = max([xlim[1], p1[0], p2[0]])
                ylim[0] = min([ylim[0], p1[1], p2[1]])
                ylim[1] = max([ylim[1], p1[1], p2[1]])
                head_width = 0.015 * np.hypot(xlim[1] - xlim[0], ylim[1] - ylim[0])
        if not head_length:
            head_length = head_width * 1.5
        r = self.ax.arrow(*p1, *dp, 
                             overhang=overhang, 
                             length_includes_head=True, 
                             linewidth=linewidth,
                             head_width=head_width,
                             head_length=head_length,
                            #  color='black',
                             zorder=self._zorder)
        self.update_axes()
        return r

    def function(self, func, x_start: float, x_end: float, num_points: int=1000, label: str=None):
        """绘制函数曲线。

        参数：
        func (function)：函数表达式，接受参数 x 并返回 y 坐标值。
        x_start (float)：x 的起始值。
        x_end (float)：x 的结束值。
        num_points (int)：绘制点的数量，默认为 1000。
        """
        self.step_zorder()
        x = np.linspace(x_start, x_end, num_points)
        y = func(x)
        r = self.ax.plot(x, y, label=label, zorder=self._zorder)
        self.update_axes()
        return r

    def parametric_equation(self, x_func, y_func, t_start, t_end, num_points=1000, label: str=None):
        """绘制参数方程。

        参数：
        x_func (function)：x(t) 的函数表达式，接受参数 t 并返回 x 坐标值。
        y_func (function)：y(t) 的函数表达式，接受参数 t 并返回 y 坐标值。
        t_start (float)：参数 t 的起始值。
        t_end (float)：参数 t 的结束值。
        num_points (int)：绘制点的数量，默认为 1000。
        """
        self.step_zorder()
        t = np.linspace(t_start, t_end, num_points)
        x = x_func(t)
        y = y_func(t)
        r = self.ax.plot(x, y, label=label, zorder=self._zorder)
        self.update_axes()
        return r

    def legend(self):
        self.step_zorder()
        self.ax.legend().set_zorder(self._zorder)
        self.update_axes()

    def save(self, filepath: str, dpi: int=300):
        """保存画布为图片。

        参数：
        filepath (str)：保存图片的路径。
        """
        self.fig.savefig(filepath, dpi=dpi)

    def clear(self):
        """清空画布。"""
        self.ax.cla()

    def close(self):
        plt.close(self.fig.number)
